#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <string>
#include <iostream>
#include <pthread.h>
#include <vector>
#include <fstream>
#include <sstream>
#include <sys/stat.h>
#include "sockets/ServerSocket.h"
#include "sockets/SocketException.h"
#include "Connection.h"
#include "Constants.h"
#include "semaphore.h"

using namespace std;

void *host_thread(void*);
void *default_listener(void*);
void *connector(void*);
void *farm_monitor(void*);
void parse_connection(string,string*,int*,string*,char*,int*,bool*);
bool checkfile(string);
int heartbeat(ServerSocket*,pthread_mutex_t*);
void cleanup();
string int2str(int);
void delete_connection(string,int);
void cleargroup(int);
void waitforworkers(int);
void parsehostfile(string);
string stripdir(string);
int str2int(string);

//For capturing ctrl+c kill, to force close properly
void sigproc(int sig);

void updatefarm();

struct host_data{
  int port;
  Connection* con;
};

const int START_PORT=30788;
const string SEPARATOR = "---"; //seperator for connection file
const int MAX_CLIENTS = 100;
const char* LOCK_FILE = "connect.lock";
const char* CONNECT_FILE = "connect.data";
const char* CONF_FILE_PRE = "config";

//current port needs to be controlled with a mutex
int current_port; 
pthread_mutex_t cpmux;
pthread_mutex_t hostargsmux;//for timing assurance

vector<Connection*> farmers;//A list of connections
pthread_mutex_t farmmux; //list mux, needed for adding or deleting connections
int overwrite_farm=0; //Ignore the current CONNECT_DATA list and just overwrite it
pthread_mutex_t owfarmmux;
string specialnumber;
sem_t specialsem;

//===================================================================//
int main(int argv, char** argc){
  cout << "server backend running...\n";
  //remove(CONNECT_FILE);
  //remove(LOCK_FILE);

  //Capture the kill signal to close properly
  signal(SIGINT,&sigproc);

  pthread_mutex_init(&cpmux,NULL);
  pthread_mutex_init(&hostargsmux,NULL);
  pthread_mutex_init(&farmmux,NULL);
  pthread_mutex_init(&owfarmmux,NULL);
  sem_init(&specialsem,0,0);

  current_port=START_PORT;

  pthread_t *hthrds;
  hthrds = (pthread_t*)malloc(sizeof(pthread_t)*MAX_CLIENTS);
  string *host_args;
  host_args = (string*)malloc(sizeof(string)*MAX_CLIENTS);

  //Start everything in reverse order, to guarantee no connections are 
  //made until server is started!
  pthread_t monitor;
  pthread_create(&monitor,NULL,farm_monitor,(void*)NULL);
  pthread_t conmake;
  pthread_create(&conmake,NULL,connector,(void*)NULL);
  pthread_t listen;
  pthread_create(&listen,NULL,default_listener,(void*)NULL);
  
  pthread_join(listen,NULL);
  pthread_join(conmake,NULL);
  pthread_join(monitor,NULL);
  
  return 0;
}

//===================================================================//
//Watches the farmers list and writes it to file from time to time
void *farm_monitor(void *args){
  printf("Started monitoring farmers...\n");
  
  while(true){
    //MASTER SLEEP STATEMENT
    sleep(10);
    updatefarm();
  }

  pthread_exit(NULL);
}

//===================================================================//
//Each client connection gets a host_thread to initialize it and control it
void *host_thread(void *args){
  printf("Spawned a stable host, connecting...\n");
  int newport,i;
  Connection* mycon; 
  string data,IP,username,command;
  string hostfile,codefile,pkfile,resultfile,numnodes;//Job data
  struct host_data *hostargs;
  ofstream pubkeyfile;
  char buffer[100];
  hostargs=(struct host_data*)args;
  newport =hostargs->port;
  mycon =hostargs->con;

  //Create the filename for the public key for this client.
  string keyfilename;

  //You're done with args, let a new host be generated.
  pthread_mutex_unlock(&hostargsmux); 

  ServerSocket datasock;
  try
    {
      // Create the socket
      ServerSocket host ( newport );
      
      host.accept(datasock);
      try
	{
	  datasock >> data;
	  if(data.compare("You there?") == 0){
	    //Correct init message recieved, 
	    //Send out response and get the IP address of the client
	    
	    data="Yessir.";
	    datasock << data;
	    
	    cout << "Stable connection to client on port " << newport << endl;
	    fflush(stdout);

	    IP=datasock.get_IP();

	    //Receive the public key file for this client
	    keyfilename=datasock.recvfile();
	    printf("Recieved pubkeyfile\n");

	    datasock>>username;
	    printf("Recieved username:%s\n",username.c_str());
	  } else {
	    printf("Bad client init message recieved in host_thread:%s\n",data.c_str());
	    pthread_exit(NULL);
	  }
	}
      catch (SocketException& e) {
	cout << "Server (host_thread) exception was caught:" << e.description() << endl;
	try{
	  datasock.close();
	  host.close();
	} catch (SocketException&){};
	pthread_exit(NULL);
      }
    }
  catch ( SocketException& e )
    {
      cout << "Server (host_thread) exception was caught:" << e.description() << "\nExiting.\n";
      pthread_exit(NULL);
    }
 
  //Close the pubkeyfile
  pubkeyfile.close();

  //Create the connection object and stick it in the vector.
  mycon->init(IP,newport);
  mycon->setpubkeyfile(keyfilename.c_str());
  mycon->setuname(username);
  pthread_mutex_lock(&farmmux);
  farmers.push_back(mycon);
  pthread_mutex_unlock(&farmmux);

  // New stable connection created, infinite loop here checking for
  // more instructions and maintaining the connection via a hearbeat
  int rsp;
  try{
    while(true){ 
      //MASTER SLEEP STATEMENT
      sleep(5);
      cout << "Switch! " << mycon->get_signature() << endl;
      switch(mycon->get_status()){

	//======================================
      case IDLE :
	//Send heartbeat
	rsp=heartbeat(&datasock,&(mycon->sendmux));
	if(rsp==TOCK){
	  printf("IDLE:Heartbeat::TickTock\n");
	  break;
	}
	if(rsp==RUNS){
	  cout << "Node " << mycon->get_signature() << " now busy\n";
	  break;
	}
	printf("NOT Received TOCK, recieved: %d\n",rsp);
	//Kill the connection
	throw SocketException ( "This connection is dead.\n");
	break;

	//======================================
      case BUSY :
	//Send heartbeat
	rsp=heartbeat(&datasock,&(mycon->sendmux));
	if(mycon->get_leader()){//MASTER
	  if(rsp==TOCK){
	    printf("BUSY:Heartbeat::MASTER-STOPPED->Kill Job\n");
	    mycon->setidle();
	    break;
	  }
	  if(rsp==RUNS){
	    //Job is running on master node
	    printf("BUSY:Heartbeat::MASTER-RUNNING\n");
	    break;
	  }
	  if(rsp==DONE){
	    printf("OMFG RECEIVING RESULTS FROM A JOB\n");
	    pthread_mutex_lock(&(mycon->sendmux));
	    //Recieve results
	    resultfile=datasock.recvfile();
	    pthread_mutex_unlock(&(mycon->sendmux));
	    command="mv "+resultfile+" "+OUTBOX;
	    system(command.c_str());
	    mycon->setfinished();
	    cleargroup(mycon->get_group());
	  }
	  if(rsp==ERR){
	    printf("BUSY:Heartbeat::MASTER-LOST\n");
	    throw SocketException ( "This connection is dead.\n");
	  }

	} else {//SLAVE
	  if(rsp==TOCK){
	    printf("BUSY:Heartbeat::SLAVE-RUNNING\n");
	    break;
	  }
	  if(rsp==ERR){
	    printf("BUSY:Heartbeat::SLAVE-LOST\n");
	    throw SocketException ( "This connection is dead.\n");
	  }
	}
	printf("NOT Received BUSY\n");
	//Kill the connection... is this right?
	mycon->kill();
	break;

	//======================================
      case NEWBUSY :
	try{
	  if(mycon->get_leader()){//Leader
	    // Prepare for a new job
    	    datasock << int2str(NBSY);
	    cout << "Sending Starting Command" << endl;
	    datasock >> data;
	    if(data.compare(int2str(ACKN))==1){
	      cout << "No Response from Client" << endl;
	      continue; // What should we do in this case?
	    }
	
	    int mygroup=mycon->get_group();
	    ifstream configfile;
	    string conf_filename=INBOX+(string)CONF_FILE_PRE;
	    conf_filename+=int2str(mygroup);
	    cout << conf_filename << endl;
	    configfile.open(conf_filename.c_str());

	    configfile.getline(buffer,100);
	    hostfile=(string)buffer;
	    
	    configfile.getline(buffer,100);
	    codefile=(string)buffer;

	    configfile.getline(buffer,100);
	    resultfile=(string)buffer;
	    
	    configfile.getline(buffer,100);
	    numnodes=(string)buffer;
	    
	    // Transfering codefile  
	    cout << "Sending Codefile" << endl;
	    pthread_mutex_lock(&(mycon->sendmux));
	    command="cp "+INBOX+codefile+" "+stripdir(codefile);
	    system(command.c_str());
	    command=stripdir(codefile);
	    datasock.sendfile(command);
	    command="rm "+stripdir(codefile);	    
	    system(command.c_str());

	    // Transfering hostfile  
	    cout << "Sending Hostfile" << endl;
	    command="cp "+INBOX+hostfile+" "+stripdir(hostfile);
	    system(command.c_str());
	    command=stripdir(hostfile);

	    //Parse the hostfile before sending, add usernames
	    parsehostfile(command);

	    datasock.sendfile(command);
	    command="rm "+stripdir(hostfile);
	    system(command.c_str());

	    // Need to send pubkeyfiles for group members to master node
	    datasock << numnodes;
	    datasock >> data;

	    // Send resultfile name
	    cout << "resultfile: " << resultfile << endl;
	    datasock << resultfile;

	    datasock >> specialnumber;
	    for(i=0;i<str2int(numnodes)-1;i++)
	      sem_post(&specialsem);

	    waitforworkers(mycon->get_group());
	    datasock << int2str(ACKN);

	    pthread_mutex_unlock(&(mycon->sendmux));
	    cout << "Job received by master client" << endl;  
	    	    
	    mycon->setbusy();
	    
	  } else {//Slave
	    pthread_mutex_lock(&(mycon->sendmux));
	    datasock << int2str(WRKR);
	    pthread_mutex_unlock(&(mycon->sendmux));
	    int mygroup=mycon->get_group();

	    pthread_mutex_lock(&farmmux);
	    vector<Connection*>::iterator itr;
	    itr=farmers.begin();
	    while(itr<farmers.end()){
	      if(mygroup==(*itr)->get_group()){
		printf("here\n");
		if(((*itr)->get_leader())){
		  //Node thats a member of the group, get its pubkey and send it 
		  pkfile=(*itr)->getpubkeyfile();
		  cout << "Sending master pubkey file to worker.\n";
		  pthread_mutex_lock(&(mycon->sendmux));
		  datasock.sendfile(pkfile);
		  pthread_mutex_unlock(&(mycon->sendmux));
		  break;
		}
	      }
	      itr++;
	    }
	    pthread_mutex_unlock(&farmmux);
	    
	    printf("waiting on special number.\n");
	    sem_wait(&specialsem);
	    printf("got the special number! %s\n",specialnumber.c_str());
	    exit(0);
	    mycon->setbusy();
	  }
	}
	catch(SocketException& e)
	  {cout << "Job Sending Error!" << endl;}  
	break;

	//======================================
      case FINISHED :
	printf("Master node finished, waiting for FE to clear.\n");
	break;

	//======================================
      case DEAD :
	printf("Killing the connection!\n");
	throw SocketException ( "This connection is dead.\n");

	//======================================
      default :
	printf("Error: Reached non-accesible state of connection.\n");
      
      }
    }
  }catch(SocketException& e){
    cout << "Heartbeat dropped: Client disconnected!" << endl;
    
    //Need to delete this connection from farmers list...
    delete_connection(mycon->get_signature(),mycon->get_group());
  }
  pthread_exit(NULL);
}

//===================================================================//
//A thread that waits for the current_port to increment creates a new host for each increment
void *connector(void *args){
  int ports=START_PORT;
  struct host_data host_args;
  Connection* newcon;
  while(true){
    //MASTER SLEEP STATMENT
    sleep(3);
    
    while(current_port>ports){
      //DEBUG
      //      printf("detected %d %d\n",current_port,ports);

      pthread_mutex_lock(&hostargsmux); //unlocked by a new host thread to ensure args are passed before being written over.
      host_args.port=ports;
      newcon=new Connection();
      host_args.con=newcon;
      pthread_t host;
      pthread_create(&host,NULL,host_thread,(void*)&host_args); 
      sleep(1);
      ports++;
    }
  }
  pthread_exit(NULL);
}

//===================================================================//
//A thread for listening on the default port and recieving connections and sending those clients new ports to connect to.
void *default_listener(void *args){
  char temp[20];
  int *args2;
  args2=(int*)args;
  try
    {
      // Create the socket
      ServerSocket server ( DEFAULT_PORT );
      
      while ( true ) {
	ServerSocket new_sock;
	
	server.accept(new_sock);
	try
	  {
	    string data;
	    new_sock >> data;
	    if(data.compare("Test message.") == 0){
	      //Correct init message recieved, 
	      //Send out new port/host data
	      pthread_mutex_unlock(&cpmux);

	      data="Message recieved.-";
	      //Stable host is current host.
	      data+=STABLE_HOST;
	      data+="-";
	      sprintf(temp,"%d",current_port);
	      data+=temp;
	      new_sock << data;
	      
	      cout << "Quick Connected a Client!\n";
	      fflush(stdout);

	      //Triggers the spawning of a new host thread/socket
	      current_port++;

	      pthread_mutex_unlock(&cpmux);
	    }	    
	  }//intial socket is destroyed and re-created in this loop
	catch ( SocketException& e ) {
	  cout << "Server exception was caught:" << e.description() << endl;
	  new_sock.close();
	  break;
	}
      }
    }
  catch ( SocketException& e )
    {
      cout << "Server exception was caught in default_listener: " << e.description() << "\nExiting.\n";
    }
  pthread_exit(NULL);
}

//Returns true if file exists
bool checkfile(string lockfilename){
  struct stat lfinfo;
  int st;
  st=stat(lockfilename.c_str(),&lfinfo);
  if(st==0)
    return true;
  return false;
}

void parse_connection(string line,string* ip,int* port,string* sig, \
		      char* stat ,int* group ,bool* leader){
  string temp;
  int loc1,loc2;

  //Get the IP
  loc1=line.find(" ");
  *ip=line.substr(0,loc1);
  //Get the Port and generate signature
  loc2=line.find(" ",loc1+1);
  temp=line.substr(loc1+1,loc2-loc1-1);
  *sig=*ip+temp;
  *port=atoi(temp.c_str());
  //Get the status
  temp=line.substr(loc2+1,1);
  *stat=(temp.c_str())[0];
  //Get the group
  loc1=line.find(" ",loc2+2);
  temp=line.substr(loc1+1,loc1-loc2+1);
  *group=atoi(temp.c_str());
  //Get the leader
  loc2=line.find(" ",loc1+1);
  temp=line.substr(loc2+1,1);
  (temp.compare("0")!=0) ? *leader=1 : *leader=0;
  //Connection parsed.
}

void updatefarm(){
  string line,signature,ip,temp;
  ifstream infile;
  ofstream lckfile,outfile;
  char buffer[100];
  int i,farmlen,cnt,found,port,group;
  char status;
  bool leader;
  vector<string> cons;
  int overwrite=0;

  if(overwrite_farm==1){
    //Do not read in the CONNECT_FILE just overwrite it
    //Useful if client dies.
    overwrite=1;
    overwrite_farm=0;
    pthread_mutex_unlock(&owfarmmux);
  }

  //Blocking statement - lock the lockfile for connections
  while(true){
    if(checkfile(LOCK_FILE)){
      // please gimme a break - Marcelo
      // bite me -Adam
      //Debug
      cout << CONNECT_FILE << " locked\n";
    }else{
      // please gimme a break - Marcelo	
      // bite me -Adam
      //Debug
      //cout << CONNECT_FILE << " unlocked\n";
      lckfile.open(LOCK_FILE);
      lckfile.close();
      break;
    }
    sleep(1);
  }
  
  //=========================
  //You own the lockfile so go ahead and work on the connection file
  cout << "Overwrite:" << overwrite << endl;
  if(checkfile(CONNECT_FILE) && !overwrite ){ //If there is a connect file read it in and parse it
    infile.open(CONNECT_FILE);
    //Read the entire connection file into cons vector
    cnt=0;
    char c;
    infile.peek();
    while(!infile.eof()){
      c=infile.peek();
      if(c=='\n'){
	infile.getline(buffer,100);
	continue;
      } 
      
      infile.getline(buffer,100);
      cons.push_back((string)buffer);
      cnt++;
      if(cnt>MAX_CLIENTS){
	cout << "Error: Too many client connections read in from " << CONNECT_FILE << " Exiting with error!\n";
	raise(SIGINT);
      }
      infile.peek();
    }
    infile.close();

    pthread_mutex_lock(&farmmux);
    farmlen=farmers.size();
    cout << "connect.data count=" << cnt << endl;
    if(cnt!=0){ //There's something in the CONNECT_FILE 
      //Go over the read in connections (done so in reverse) 
      //and modify local connections accordingly
      while(cons.size()>0){ //go over altered connections
	line=cons.back();
	cons.pop_back();
	if(line.compare(SEPARATOR)!=1){ //break to new connections 
	  break;
	}
	//Parse the connection info
	parse_connection(line,&ip,&port,&signature,&status,&group,&leader);
	
	found=0;
	for(i=0;i<farmlen;i++){
	  temp=farmers[i]->get_signature();
	  //cout << temp << " " << signature << endl;
	  if(temp.compare(signature)==0){
	    found=1;
	    break;
	  }
	}
	if(found==1){
	  //Having read in the status of the farmers:
	  //Decide what to do based on the status change
	  //cout << "FOUND!\n";
	  switch(status)
	    {
	    case BUSY:
	      //Nothing new to do.
	      break;

	    case NEWBUSY:
	      if(farmers[i]->get_status()==IDLE){
		printf("group:%i,leader:%i\n",group,leader);
		farmers[i]->setnewbusy(group,leader);
		break;
	      }
	      if(farmers[i]->get_status()==NEWBUSY){
		//This connection is still setting up. Don't do anything.
		break;
	      }
	      cout << "ERROR: trying to change farmer (" << signature << ") to state newbusy when it was already in state: " << farmers[i]->get_status() << " which is not idle.\n";
	      //Should probably reset the connection at this point?
	      break;

	    case IDLE:
	      farmers[i]->setidle();
	      break;

	    case DEAD:
	      farmers[i]->kill();
	      break;

	    case FINISHED:
	      //Nothing new to do.
	      break;

	    default:
	      cout << "ERROR: invalid status setting for farmer (" << ip << ")\n";
	      break;
	    }
	}
	else{
	  cout << "WARNING!\n";
	  cout << "Invalid connection written by front end to " << CONNECT_FILE << ".\nIgnoring connection this connection. FIX THIS ERROR!\n";
	}
      }
    }
  }
  //Now write the modified list to the file
  outfile.open(CONNECT_FILE);
  for(i=0;i<farmlen;i++){
    farmers[i]->print(&outfile);
    //Debug
    farmers[i]->print_scr();
  }
  outfile.close();
  
  pthread_mutex_unlock(&farmmux); 
  if(overwrite)
    printf("Forced loop over connections.\n");
  else
    printf("looped over connections.\n");
  //=========================
  
  //Unlock the lockfile    
  remove(LOCK_FILE);

  return;
}

void sigproc(int sig){ 	
  if(sig==SIGINT){
    cout << "\nCleaning up...\n";
    cleanup();
    cout << "Closing.\n";
    exit(0);
  }
}

int heartbeat(ServerSocket* sock,pthread_mutex_t* smux){
  string data;
  pthread_mutex_lock(smux);
  try{
    *sock<<int2str(TICK);
    *sock>>data;
  } catch(SocketException& e) {
    cout << "Server (host_thread) exception was caught:" << e.description() << endl;
    cout << "Heartbeat Error!" << endl;
    return ERR;
  }
  pthread_mutex_unlock(smux);
  return atoi(data.c_str());
}

void cleanup(){
  //Remove Operating files
  remove(LOCK_FILE);
  remove(CONNECT_FILE);

  //Remove the keyfiles
  //system("rm -rf keys/*");

  //Remove any user data files
}

string int2str(int i){
  stringstream out;
  out << i;
  return out.str();
}

//Remove the connection with the signature given from the farm and make sure
//the socket is closed.
void delete_connection(string signature,int group){

  pthread_mutex_lock(&farmmux);
  vector<Connection*>::iterator itr;
  itr=farmers.begin();
  string remove_pk;
  while(itr<farmers.end()){
     //(*itr)->print_scr();

    //Set all busy clients in same group to not busy.
    if(group==(*itr)->get_group()){
      (*itr)->setidle();
    }

    //Remove all the dead connection
    if(signature.compare((*itr)->get_signature())==0){
      remove((*itr)->getpubkeyfile());
      //system(remove_pk.c_str());
      farmers.erase(itr);
    }
    itr++;
  }
  pthread_mutex_unlock(&farmmux);
  
  //Force the current state to be written in CONNECT_FILE
  pthread_mutex_lock(&owfarmmux);
  overwrite_farm=1;
  updatefarm();
}

void cleargroup(int group){
  //Go over all connections and if they are in this group set them to idle (non-leaders)
  pthread_mutex_lock(&farmmux);
  vector<Connection*>::iterator itr;
  itr=farmers.begin();
  while(itr<farmers.end()){
    //Set all busy clients in same group to not busy.
    if(group==(*itr)->get_group()){
      if(!(*itr)->get_leader())
	(*itr)->setidle();
    }
    itr++;
  }
  pthread_mutex_unlock(&farmmux);
}

void waitforworkers(int group){
  //Go over all connections and if they are in this group wait for them to be set to busy
  /*int allbusy=1;
  while(true){
    pthread_mutex_lock(&farmmux);
    vector<Connection*>::iterator itr;
    itr=farmers.begin();
    while(itr<farmers.end()){
      //Set all busy clients in same group to not busy.
      if(group==(*itr)->get_group()){
	printf("hiohoihoiasdf\n");
	if(!(*itr)->get_leader())
	  if((*itr)->get_status()!=BUSY){
	    allbusy=0;
	    break;
	  }
      }
      itr++;
    }
    pthread_mutex_unlock(&farmmux);
    if(allbusy){
      break;
      printf("All Busy!\n");
    }
    sleep(1);
    }*/
  sleep(5);
}

void parsehostfile(string hostfile){
  ifstream hfile;
  hfile.open(hostfile.c_str());
  string ip;
  char buff[128];
  pthread_mutex_lock(&farmmux);
  vector<Connection*>::iterator itr;
  vector<string> ipuns;
  hfile.getline(buff,128);
  while(!hfile.eof()){
    ip=buff;
    itr=farmers.begin();
    while(itr<farmers.end()){
      //Set all busy clients in same group to not busy.
      if(((*itr)->get_ip()).compare(ip)==0){
	ipuns.push_back((*itr)->get_uname()+"@"+ip);
	break;
      }
      itr++;
    }
    hfile.getline(buff,128);
  }

  pthread_mutex_unlock(&farmmux);
  
  hfile.close();
  ofstream nhfile;
  nhfile.open(hostfile.c_str());
  for(int i=0;i<(int)ipuns.size();i++){
    nhfile<< (string)ipuns[i];
    nhfile<< "\n";
  }
  nhfile.close();
}

string stripdir(string path){
  int loc;
  loc=path.rfind("/");
  return path.substr(loc+1,path.size());
}

int str2int(string s){
  return atoi(s.c_str());
}
